from setuptools import setup, find_packages
from torch.utils.cpp_extension import CppExtension, BuildExtension
import os
import glob

include_dirs = os.path.dirname(os.path.abspath(__file__))
print(include_dirs)

ext_modules=[
    CppExtension(
        name="conv_extension",
        language="c++",
        sources=["./con_batch_act.cpp"],  
        include_dirs=[include_dirs,
                        '/usr/local/cuda-10.0/include',
                        '/data/zh/anaconda3/lib/python3.6/site-packages/torch/lib/include',
                        '/data/zh/anaconda3/lib/python3.6/site-packages/torch/lib/include/torch/csrc/api/include',
                        '/data/zh/anaconda3/include/python3.6m'],
        library_dirs=['/usr/local/cuda-10.0/lib64'],
        libraries = ['cudnn', 'cudart']
    )
]

setup(
    name="conv_extension", 
    version = "1.0",
    packages=find_packages(),
    ext_modules=ext_modules,
    cmdclass={'build_ext': BuildExtension}
)



